import torch
import torch.nn as nn
import torch.nn.functional as F

from typing import Any, Dict
from torchvision.models import resnet101
from torch.hub import load_state_dict_from_url
from torchvision.models._utils import IntermediateLayerGetter


class AtrousSpatialPyramidPooling(nn.Module):
    def __init__(self, in_channels, dilation_rate=[12, 24, 36]):
        super(AtrousSpatialPyramidPooling, self).__init__()
        self.out_1x1 = nn.Sequential(
            nn.Conv2d(in_channels, 256, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )

        self.out_3x3_1 = self._make_block(in_channels, dilation_rate=dilation_rate[0])
        self.out_3x3_2 = self._make_block(in_channels, dilation_rate=dilation_rate[1])
        self.out_3x3_3 = self._make_block(in_channels, dilation_rate=dilation_rate[2])

        self.out_image_pooling_conv = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, 256, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )

        self.final_conv = nn.Sequential(
            nn.Conv2d(1280, 256, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5)
        )

    def _make_block(self, in_channels, dilation_rate):
        return nn.Sequential(
            nn.Conv2d(in_channels, 256, kernel_size=3, stride=1, padding=dilation_rate, dilation=dilation_rate, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        out_1x1 = self.out_1x1(x)

        out_3x3_1 = self.out_3x3_1(x)
        out_3x3_2 = self.out_3x3_2(x)
        out_3x3_3 = self.out_3x3_3(x)

        out_image_pooling = self.out_image_pooling_conv(x)
        out_image_pooling = F.interpolate(out_image_pooling, size=x.size()[2:], mode='bilinear', align_corners=False)

        x = torch.cat([out_1x1, out_3x3_1, out_3x3_2, out_3x3_3, out_image_pooling], dim=1)

        x = self.final_conv(x)

        return x


class ResNet101Backbone(nn.Module):
    def __init__(self):
        super(ResNet101Backbone, self).__init__()
        
        backbone = resnet101(pretrained=False, replace_stride_with_dilation=[False, False, True])
        state_dict = load_state_dict_from_url('https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', progress=True)
        backbone.load_state_dict(state_dict, strict=False)
            
        self.backbone = IntermediateLayerGetter(
            backbone,
            return_layers={'layer4': 'out', 'layer1': 'low_level'}
        )

    def forward(self, x):
        out = self.backbone(x)
        return out['out'], out['low_level']


class DeeplabV3Plus(nn.Module):
    def __init__(self, cfg: Dict[str, Any]):
        super(DeeplabV3Plus, self).__init__()
        num_classes = cfg.get('num_classes')

        self.backbone = ResNet101Backbone()
        self.aspp = AtrousSpatialPyramidPooling(2048)

        self.low_level_conv = nn.Conv2d(256, 48, kernel_size=1, bias=False)
        self.low_level_bn = nn.BatchNorm2d(48)

        self.final_conv1 = nn.Conv2d(304, 256, kernel_size=3, padding=1, bias=False)
        self.final_bn1 = nn.BatchNorm2d(256)
        self.final_conv2 = nn.Conv2d(256, num_classes, kernel_size=1)

    def forward(self, x):
        size = x.size()[2:]

        x, low_level_features = self.backbone(x) # 2048, 256
        x = self.aspp(x) # 256
        x = F.interpolate(x, size=low_level_features.shape[-2:], mode='bilinear', align_corners=False)

        low_level_features = F.relu(self.low_level_bn(self.low_level_conv(low_level_features)))

        x = torch.cat([x, low_level_features], dim=1)

        x = F.relu(self.final_bn1(self.final_conv1(x)))
        x = F.interpolate(x, size=size, mode='bilinear', align_corners=False)
        x = self.final_conv2(x)

        return {'out': x}
